# *-* coding:utf8 *-*

import sys
from PyQt5.QtWidgets import QMainWindow ,QApplication,QMessageBox

from PyQt5.Qt import QThread,QMutex

import gym
from stable_baselines3 import PPO
from torch._C import set_flush_denormal
import puck_world
from train import Ui_Form as train_window
import threading
import time
qmut_1 = QMutex() # 创建线程锁
qmut_2 = QMutex()
class Thread_expert(QThread): 
    def __init__(self,ui,turn,model_save_path):

        super().__init__()
        self.ui = ui
        self.stop_train_expert_flag = 1
        self.turn = turn
        self.model_save_path = model_save_path
    def stop(self):
        self.stop_train_expert_flag = 0
        self.requestInterruption()
        self.wait()
    def run(self):
        qmut_1.lock() # 加锁
        self.train()
        qmut_1.unlock() # 解锁
    def train(self):
        env = gym.make("puck-world-v0")
        model = PPO("MlpPolicy", env, verbose=1)
        time_all = time.time()
        for update in range(self.turn):
            time_start = time.time()
            model.learn(total_timesteps=1000)
            model.save(self.model_save_path+"第"+str(update)+"代")
            time_end = time.time()
            self.ui.textBrowser.append("已训练代数"+str(update)+"代，距离结束还有"+str(self.turn-update)+"代")
            time_cost_all = str(time_end -time_all)
            time_cost = str((time_end-time_start)*(self.turn-update))
            message_time = "总耗时"+time_cost_all+"秒，距离结束还有"+time_cost+"秒"
            self.ui.textBrowser.append(message_time)
            
            if not self.stop_train_expert_flag:
                message = "训练已停止"
                self.ui.textBrowser.append(message)
                time.sleep(2)
                env.close()
                return 0
        env.close()

class Thread_imitation(QThread): 
    def __init__(self,ui,turn,model_save_path,expert = None):
        super().__init__()
        self.ui = ui
        self.stop_train_flag = 1
        self.turn = turn
        self.model_save_path = model_save_path
    def stop(self):
        self.stop_train_expert_flag = 0
        self.requestInterruption()
        self.wait()
    def run(self):
        qmut_2.lock() # 加锁
        self.train()
        qmut_2.unlock() # 解锁
    def train(self):
        env = gym.make("puck-world-v0")
        model = PPO("MlpPolicy", env, verbose=1)
        time_all = time.time()
        for update in range(self.turn):
            time_start = time.time()
            model.learn(total_timesteps=1000)
            model.save(self.model_save_path+"第"+str(update)+"代")
            time_end = time.time()
            self.ui.textBrowser_2.append("已训练代数"+str(update)+"代，距离结束还有"+str(self.turn-update)+"代")
            time_cost_all = str(time_end -time_all)
            time_cost = str((time_end-time_start)*(self.turn-update))
            message_time = "总耗时"+time_cost_all+"秒，距离结束还有"+time_cost+"秒"
            self.ui.textBrowser_2.append(message_time)
            if not self.stop_train_flag:
                message = "训练已停止"
                self.ui.textBrowser_2.append(message)
                time.sleep(2)
                env.close()
                return 0
        env.close()


class MainWindow(QMainWindow):
    def __init__(self,parent=None):
        super(MainWindow,self).__init__(parent)
        self.train_expert_flag = 0
        self.stop_train_expert_flag =1
        self.stop_train_imitation_flag=0
        self.train_imitation_flag = 0
        self.ui = train_window()
        self.ui.setupUi(self)
        self.setWindowTitle("模型训练")

        self.connect_reinforce_train_button()
        self.connect_imitation_train_button()
        self.connect_reinforce_stop_button()
        self.connect_imitation_stop_button()
    def train_expert(self):
        if self.ui.lineEdit_2.text():
            model_save_path = self.ui.lineEdit_2.text()
        else:
            model_save_path="/home/maxwene/Project/mov_qt_ui/model/model"
        if self.ui.lineEdit.text():
            turn = int(self.ui.lineEdit.text())
        else:
            turn = 500
        return turn,model_save_path

    def train_imitation(self):
        if self.ui.lineEdit_3.text():
            model_save_path = self.ui.lineEdit_5.text()
        else:
            model_save_path="/home/maxwene/Project/mov_qt_ui/model/imitation_model"

        if self.ui.lineEdit_3.text():
            turn = int(self.ui.lineEdit.text())
        else:
            turn = 500
        expert_path = 0
        return turn,model_save_path


    def stop_expert_train(self):
        if self.thread_expert:
            self.thread_expert.stop_train_expert_flag = 0

    def stop_imitation_train(self):
        if self.thread_imitation:
            self.thread_imitation.stop_train_flag = 0

    def train_expert_thread(self):
        turn ,model_save_path = self.train_expert()
        self.thread_expert = Thread_expert(ui = self.ui,turn=turn,model_save_path=model_save_path)  # 创建线程
        self.thread_expert.start()  # 开始线程



    def train_imitation_thread(self):
        turn ,model_save_path = self.train_imitation()
        self.thread_imitation = Thread_imitation(ui = self.ui,turn=turn,model_save_path=model_save_path)  # 创建线程
        self.thread_imitation.start()  # 开始线程

    def connect_reinforce_train_button(self):
        self.ui.pushButton.clicked.connect(self.train_expert_thread)
    def connect_imitation_train_button(self):
        self.ui.pushButton_2.clicked.connect(self.train_imitation_thread)
    def connect_reinforce_stop_button(self):
        self.ui.pushButton_3.clicked.connect(self.stop_expert_train)
    def connect_imitation_stop_button(self):
        self.ui.pushButton_4.clicked.connect(self.stop_imitation_train)
if __name__ =="__main__":
    app = QApplication(sys.argv)
    B = MainWindow()
    B.show()
    sys.exit(app.exec_())